package models

import (
	"fmt"
	"log"
	"math"
	"strings"
	"sync"
)

// EmbeddingModel 定义嵌入模型接口
type EmbeddingModel interface {
	// 初始化模型
	Initialize(config *ModelConfig) error
	// 关闭模型
	Close() error
	// 获取文本向量
	GetEmbedding(text string) ([]float32, error)
	// 批量获取文本向量
	GetEmbeddings(texts []string) ([][]float32, error)
	// 获取向量维度
	GetDimension() int
	// 获取模型信息
	GetModelInfo() *ModelInfo
}

// ModelConfig 模型配置
type ModelConfig struct {
	ModelName  string `json:"model_name"`
	ModelPath  string `json:"model_path"`
	MaxLength  int    `json:"max_length"`
	BatchSize  int    `json:"batch_size"`
	UseGPU     bool   `json:"use_gpu"`
	NumThreads int    `json:"num_threads"`
	CacheSize  int    `json:"cache_size"`
}

// ModelInfo 模型信息
type ModelInfo struct {
	Name        string `json:"name"`
	Version     string `json:"version"`
	Dimension   int    `json:"dimension"`
	MaxLength   int    `json:"max_length"`
	Language    string `json:"language"`
	Description string `json:"description"`
}

// DefaultModelConfig 创建默认模型配置
func DefaultModelConfig() *ModelConfig {
	return &ModelConfig{
		ModelName:  "bge-base-zh-v1.5",
		ModelPath:  "./models/bge-base-zh-v1.5",
		MaxLength:  512,
		BatchSize:  32,
		UseGPU:     false,
		NumThreads: 4,
		CacheSize:  1000,
	}
}

// ChineseEmbeddingModel 中文嵌入模型实现
type ChineseEmbeddingModel struct {
	config    *ModelConfig
	modelInfo *ModelInfo
	cache     *EmbeddingCache
	mutex     sync.RWMutex
	loaded    bool
}

// NewChineseEmbeddingModel 创建中文嵌入模型
func NewChineseEmbeddingModel() *ChineseEmbeddingModel {
	return &ChineseEmbeddingModel{
		config: DefaultModelConfig(),
		modelInfo: &ModelInfo{
			Name:        "bge-base-zh-v1.5",
			Version:     "1.5",
			Dimension:   768,
			MaxLength:   512,
			Language:    "zh",
			Description: "BGE中文向量模型，支持中文语义理解和向量化",
		},
		cache: NewEmbeddingCache(1000),
	}
}

// Initialize 初始化模型
func (m *ChineseEmbeddingModel) Initialize(config *ModelConfig) error {
	m.mutex.Lock()
	defer m.mutex.Unlock()

	if config != nil {
		m.config = config
	}

	log.Printf("正在初始化中文嵌入模型: %s", m.config.ModelName)

	// TODO: 实际的模型加载逻辑
	// 这里应该加载实际的BGE模型文件
	// 目前使用模拟实现

	m.loaded = true
	log.Printf("中文嵌入模型初始化完成，维度: %d", m.modelInfo.Dimension)

	return nil
}

// Close 关闭模型
func (m *ChineseEmbeddingModel) Close() error {
	m.mutex.Lock()
	defer m.mutex.Unlock()

	if !m.loaded {
		return nil
	}

	log.Println("正在关闭中文嵌入模型...")

	// 清理缓存
	if m.cache != nil {
		m.cache.Clear()
	}

	m.loaded = false
	log.Println("中文嵌入模型已关闭")

	return nil
}

// GetEmbedding 获取单个文本的向量
func (m *ChineseEmbeddingModel) GetEmbedding(text string) ([]float32, error) {
	m.mutex.RLock()
	defer m.mutex.RUnlock()

	if !m.loaded {
		return nil, fmt.Errorf("模型未初始化")
	}

	if text == "" {
		return nil, fmt.Errorf("输入文本不能为空")
	}

	// 检查缓存
	if cached := m.cache.Get(text); cached != nil {
		return cached, nil
	}

	// 预处理文本
	processedText := m.preprocessText(text)

	// 生成向量（模拟实现）
	vector := m.generateEmbedding(processedText)

	// 缓存结果
	m.cache.Set(text, vector)

	return vector, nil
}

// GetEmbeddings 批量获取文本向量
func (m *ChineseEmbeddingModel) GetEmbeddings(texts []string) ([][]float32, error) {
	if len(texts) == 0 {
		return nil, fmt.Errorf("输入文本列表不能为空")
	}

	results := make([][]float32, len(texts))

	// 分批处理
	batchSize := m.config.BatchSize
	for i := 0; i < len(texts); i += batchSize {
		end := i + batchSize
		if end > len(texts) {
			end = len(texts)
		}

		batch := texts[i:end]
		for j, text := range batch {
			vector, err := m.GetEmbedding(text)
			if err != nil {
				return nil, fmt.Errorf("处理文本 %d 失败: %v", i+j, err)
			}
			results[i+j] = vector
		}
	}

	return results, nil
}

// GetDimension 获取向量维度
func (m *ChineseEmbeddingModel) GetDimension() int {
	return m.modelInfo.Dimension
}

// GetModelInfo 获取模型信息
func (m *ChineseEmbeddingModel) GetModelInfo() *ModelInfo {
	return m.modelInfo
}

// preprocessText 预处理文本
func (m *ChineseEmbeddingModel) preprocessText(text string) string {
	// 去除多余空白字符
	text = strings.TrimSpace(text)

	// 限制文本长度
	if len([]rune(text)) > m.config.MaxLength {
		runes := []rune(text)
		text = string(runes[:m.config.MaxLength])
	}

	return text
}

// generateEmbedding 生成向量（模拟实现）
func (m *ChineseEmbeddingModel) generateEmbedding(text string) []float32 {
	// 这是一个模拟实现，实际应该调用真实的BGE模型
	dimension := m.modelInfo.Dimension
	vector := make([]float32, dimension)

	// 基于文本内容生成确定性的向量
	hash := m.simpleHash(text)

	for i := 0; i < dimension; i++ {
		// 使用简单的伪随机生成，确保相同文本产生相同向量
		seed := hash + uint32(i)
		value := float32(seed%1000)/1000.0 - 0.5 // 范围 [-0.5, 0.5]
		vector[i] = value
	}

	// 归一化向量
	m.normalizeVector(vector)

	return vector
}

// simpleHash 简单哈希函数
func (m *ChineseEmbeddingModel) simpleHash(text string) uint32 {
	var hash uint32 = 5381
	for _, char := range text {
		hash = ((hash << 5) + hash) + uint32(char)
	}
	return hash
}

// normalizeVector 向量归一化
func (m *ChineseEmbeddingModel) normalizeVector(vector []float32) {
	var norm float32
	for _, v := range vector {
		norm += v * v
	}
	norm = float32(math.Sqrt(float64(norm)))

	if norm > 0 {
		for i := range vector {
			vector[i] /= norm
		}
	}
}

// EmbeddingCache 向量缓存
type EmbeddingCache struct {
	cache   map[string][]float32
	maxSize int
	mutex   sync.RWMutex
}

// NewEmbeddingCache 创建向量缓存
func NewEmbeddingCache(maxSize int) *EmbeddingCache {
	return &EmbeddingCache{
		cache:   make(map[string][]float32),
		maxSize: maxSize,
	}
}

// Get 获取缓存的向量
func (c *EmbeddingCache) Get(key string) []float32 {
	c.mutex.RLock()
	defer c.mutex.RUnlock()

	if vector, exists := c.cache[key]; exists {
		// 返回副本以避免并发修改
		result := make([]float32, len(vector))
		copy(result, vector)
		return result
	}

	return nil
}

// Set 设置缓存
func (c *EmbeddingCache) Set(key string, vector []float32) {
	c.mutex.Lock()
	defer c.mutex.Unlock()

	// 如果缓存已满，清理一些旧条目
	if len(c.cache) >= c.maxSize {
		// 简单的清理策略：清理一半
		count := 0
		for k := range c.cache {
			delete(c.cache, k)
			count++
			if count >= c.maxSize/2 {
				break
			}
		}
	}

	// 存储副本以避免外部修改
	cached := make([]float32, len(vector))
	copy(cached, vector)
	c.cache[key] = cached
}

// Clear 清空缓存
func (c *EmbeddingCache) Clear() {
	c.mutex.Lock()
	defer c.mutex.Unlock()

	c.cache = make(map[string][]float32)
}

// Size 获取缓存大小
func (c *EmbeddingCache) Size() int {
	c.mutex.RLock()
	defer c.mutex.RUnlock()

	return len(c.cache)
}
